#!/usr/bin/python

import roslib; roslib.load_manifest('efri') 
import rospy
import m3.toolbox as m3t
import sys, time, os
import math, numpy as np
import time
import hrl_segway_omni.segway as segway
import cPickle as pkl

import zenither.zenither as zen
import thread
import hrl_cody_arms.cody_arm_client as cac
import equilibrium_point_control.epc as epc
import hrl_lib.util as ut
import rospy
import traceback
from optparse import OptionParser
from hrl_msgs.msg import FloatArray
from std_msgs.msg import Float64

##
#Class Lead()
#Allows a user to grab the left/right/both MEKA arms in order to control Segway RMP motion (x).
#
class Lead():
    ##
    #Constructor for Lead() class
    #@param rob an hrl_robot.M3HrlRobot() object.
    #@param arms a string containing one of: 'left_arm','right_arm', or 'both' to denote which arms to control (left or right or both arms). It is assumed that the arms you select to control were initialized in the "rob" object that is also passed to this constructor.
    #@param mec a segway.Mecanum() object. 
    #@param z a Zenither HRL2 object
    #@param log_name - file to log the data to (default is /dev/null: no file created).
    def __init__(self,l_arm,r_arm,arms,mec,z,gain,log_name='/dev/null',max_allowable_height=None,
                 init_height=None):
        self.out_log = open(log_name,'w')
        self.lac = l_arm
        self.rac = r_arm 
        self.l_epc = epc.EPC(l_arm)
        self.r_epc = epc.EPC(r_arm)
        self.gain = gain
        self.arms = arms
        self.mec = mec
        self.z = z

        self.run_thread_flag = False
        self.thread_stopped_flag = True
    
    def run(self,q0_left,q0_right,xoffset):
        self.force_l_rot_pub = rospy.Publisher('efri/l_arm/force_rot', FloatArray )
        self.xvel_pub = rospy.Publisher('efri/xvel', FloatArray )
        self.run_thread_flag = True
        self.thread_stopped_flag = False

        
        print '------------------------'
        if self.arms == 'left_arm' or self.arms == 'both':
            self.l_epc.go_jep(q0_left)
            time.sleep(2)
            self.lac.bias_wrist_ft()
            print >> self.out_log, 'Force bias for left arm:'
            print >> self.out_log, self.lac.fts_bias['force']
        if self.arms == 'right_arm' or self.arms == 'both':
            self.r_epc.go_jep(q0_right)
            time.sleep(2)
            self.rac.bias_wrist_ft()
            print >> self.out_log, 'Force bias for right arm:'
            print >> self.out_log, self.rac.fts_bias['force']
        if self.arms == 'left_arm' or self.arms == 'both':
            q0_left = self.lac.get_joint_angles()
            p = self.lac.kinematics.FK(q0_left)
            p = p[0]    #grab the first element of the list (position)
            print >>self.out_log, 'home position angles are:\n (left): ', q0_left
            print >>self.out_log, 'home positions are (left):\n', p
        if self.arms == 'right_arm' or self.arms == 'both':
            q0_right = self.rac.get_joint_angles()
            p_right = self.rac.kinematics.FK(q0_right)
            p_right = p_right[0]
            print >>self.out_log, 'home position angles are:\n(right): ',q0_right
            print >>self.out_log, 'home positions are (right):\n', p_right
    
        max_xvel = 0.70  #m/s
        x_thres = 1.0
        y_thres = 2.0
        max_human_xforce = 15 #N
        
        if self.gain == 'h':
            #high gain setting
            x_velscale = max_xvel/max_human_xforce/2.0 #0.02 m/sN
        elif self.gain == 'l':
            #low gain setting
            x_velscale = max_xvel/max_human_xforce/4.0 #0.01 m/sN
        global y_velscale
            
        lastcmd = time.time()
        
     
        len = 3 #number of samples to save in history.
        xvel_history = np.matrix(np.zeros([len,1]))
        fx_history = np.matrix(np.zeros([len,1]))
        print 'The robot is now READY to be lead around.'
        print 'Hit a key to START trial timing.'
        print 'waiting for keystroke...'

        try:
            while (self.run_thread_flag):
                collision_detected = False #reset before each trial
                
                if self.lac.is_motor_power_on() == False or self.rac.is_motor_power_on() == False:
                    break 
                
                if self.arms == 'left_arm' or self.arms == 'both':
                    force_left = FloatArray()
                    force_left_rot = FloatArray()
                    [force_left.data, rot] = self.lac.get_wrist_force(base_frame=True,filtered=True,return_transform=True) 
                    force_left_rot.data = (np.reshape(rot,9, order='F')).T
                    fx_left = force_left.data[0]
                    fy_left = force_left.data[1]
                else:
                    fx_left = 0.0


                if self.arms == 'right_arm' or self.arms == 'both':
                    force_right = FloatArray()
                    force_right_rot = FloatArray()
                    force_right.data = self.rac.get_wrist_force(base_frame=True) 
                    fx_right = force_right.data[0]
                    fy_right = force_right.data[1]
                else:
                    fx_right = 0.0
                
                fx = fx_left + fx_right 
                fy = 0.0                

                fx_history[0:(len-1)] = xvel_history[1:len]
                fx_history[(len-1)] = fx
                
                #Determine the sign (+/-) of control parameters
                xsign = np.sign(fx)
                
                #Compute x-vel
                if (xsign == -np.sign(fx_history[len-2]) and np.abs(fx-fx_history[len-2]) > 10):  
                    collision_detected = True

                if (fx == 0 and fy == 0): 
                    xvel = 0.0
                else:
                    xvel = xsign*np.min([x_velscale*np.float(np.abs(fx)), max_xvel])

            
                xvel_history[0:(len-1)] = xvel_history[1:len]
                xvel_history[(len-1)] = xvel      #take into account current computed xvel
                xvel = np.mean(xvel_history.tolist())   #replace xvel with the mean of the history
                xvel_history[(len-1)] = xvel  #save the actual commanded xvel
                #Set velocity of mobile base. Set y-vel and a-vel to zero.
	        self.mec.set_velocity(xvel+xoffset,0.0,0.0)  
                
                commanded_xvel = FloatArray()
                commanded_xvel.data = [xvel]

                curr_time = rospy.Time.now()
                force_left_rot.header.stamp = curr_time
                force_left_rot.header.stamp = curr_time
                commanded_xvel.header.stamp = curr_time

                self.force_l_rot_pub.publish(force_left_rot)
                self.xvel_pub.publish(commanded_xvel)

                rospy.sleep(.05)

            self.thread_stopped_flag = True
            self.stop() 
        except:  
            self.thread_stopped_flag = True
            (type, value, tb) = sys.exc_info()
            print type
            print value
            print tb
            follower.stop()
            raise
    
    
    def start_lead_thread(self,ql=None,qr=None,xoffset=0.0):
        if ql == None:
            ql = [0, 0, 0, math.pi/2, 0, 0, 0]  
        if qr == None:
            qr = [0, 0, 0, math.pi/2, 0, 0, 0]  

        thread.start_new_thread(self.run,(ql,qr,xoffset))

    def stop(self):
        self.run_thread_flag = False
        while self.thread_stopped_flag == False:
            time.sleep(0.005) # wait for the thread to end.
        print 'engaging zenither lock...'
        try:
            self.out_log.close()
        except (AttributeError):
            pass
            

if __name__ == '__main__':

    parser = OptionParser()
    parser.add_option("-n",type="float",dest="xoffset")
    parser.add_option("-p","--participant",type="string",dest="participant_num")
    parser.add_option("-t","--trial",type="string",dest="trial")
    parser.add_option("-g","--gain",type="string",dest="gain")
    parser.add_option("-s","--stiffness",type="string",dest="stiffness")

    (options, args) = parser.parse_args()
    if (options.participant_num == None or options.trial == None or
    options.gain == None or options.stiffness == None):
        print 'You must specify participant number (-p), trial number (-t), gain (-g), and stiffness (-s) for each trial. Try again.'
        sys.exit()
    else:
        name = 'sub' + options.participant_num + '_trial_' + options.trial + '_gain_' + options.gain + '_stiff_' + options.stiffness
        print '*Results will be saved to: ', name

    if options.xoffset == None:
        xoffset = 0.0
    else:
        xoffset = options.xoffset
    print '*Cody is leading with xoffset of ', xoffset , ' m/s (default is 0.0 m/s)'
    base = segway.Mecanum()    

    zed = None

    rac = cac.CodyArmClient('r')
    lac = cac.CodyArmClient('l')
    rac.kinematics.set_tooltip(p=np.matrix([0.,0.,-.08895]).T) #The centerline of the rubber balls is .08895m from the flat surface of the last link plate
    lac.kinematics.set_tooltip(p=np.matrix([0.,0.,-.08895]).T)
    
    #Set stiffness settings
    if options.stiffness == 'l':
        stiffness_list_r = list([0.319,0.3,0.8,0.598,0.1])
        stiffness_list_l = list([0.009,0.3,0.8,0.137,0.1])    
        q0_right = [-0.029671, 0, 0, 1.864012, 0, 0, 0]  
        q0_left = [0.0, 0, 0, 1.7, 0, 0, 0]  
    elif options.stiffness == 'h': 
        stiffness_list_r = list([1.0,0.3,0.8,1.0,0.1])      
        stiffness_list_l = list([1.0,0.3,0.8,1.0,0.1])     
        q0_left = [0.083776, 0, 0, 1.740093, 0, 0, 0]  
        q0_right = [0, 0, 0, 1.720895, 0, 0, 0]         
        

    rac.set_impedance_scale(stiffness_list_r)
    lac.set_impedance_scale(stiffness_list_l)
    follower = Lead(lac, rac,'both',base,zed,options.gain,log_name=(name+'.log'))

    print '------------------------'
    print '\n\n\n1. Tell the ESTOP person to pull up the button. '
    print '\n2. Hit any key to power up the arms.'
    print 'waiting for keystroke...\n\n\n\n\n'
    k=m3t.get_keystroke()
    follower.start_lead_thread(q0_left,q0_right,xoffset)

    try:
        k=m3t.get_keystroke()
        print 'Hit "Control-C" to END trial timing.'
        print 'waiting for "Control-C"...'
        
        while not rospy.is_shutdown():
            continue

        follower.stop()
    except:  
        follower.stop()
        print 'stopping MEKA arms...'
        raise

